use futures_util::{SinkExt, StreamExt};
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::sync::Mutex;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use uuid::Uuid;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let url = "ws://127.0.0.1:8080";
    
    println!("Connecting to WebSocket server at: {}", url);
    
    let (ws_stream, _) = connect_async(url).await?;
    println!("WebSocket handshake completed");
    
    let (mut write, mut read) = ws_stream.split();
    
    // Generate user ID and username
    let user_id = Uuid::new_v4();
    let username = format!("user_{}", user_id.to_string()[..8].to_string());
    
    // Send join message
    let join_msg = json!({
        "type": "join",
        "user_id": user_id,
        "username": username
    });
    
    let join_msg_str = join_msg.to_string();
    println!("Sending join message: {}", join_msg_str);
    write.send(Message::Text(join_msg_str)).await?;
    println!("Joined as: {}", username);
    println!("Your user ID: {}", user_id);
    
    let user_id_clone = user_id;
    let online_users = Arc::new(Mutex::new(HashMap::new())); // user_id -> username
    
    // Clone Arc for the read task
    let online_users_clone = Arc::clone(&online_users);
    
    // Spawn task to handle incoming messages
    let read_task = tokio::spawn(async move {
        while let Some(msg) = read.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&text) {
                        match parsed["type"].as_str() {
                            Some("broadcast") => {
                                let sender = parsed["sender_id"].as_str().unwrap_or("unknown");
                                let content = parsed["content"].as_str().unwrap_or("");
                                println!("[BROADCAST from {}] {}", sender, content);
                            }
                            Some("direct") => {
                                let sender = parsed["sender_id"].as_str().unwrap_or("unknown");
                                let content = parsed["content"].as_str().unwrap_or("");
                                println!("[DIRECT from {}] {}", sender, content);
                            }
                            Some("join") => {
                                let user = parsed["username"].as_str().unwrap_or("unknown");
                                let user_id = parsed["user_id"].as_str().unwrap_or("unknown");
                                println!("[SYSTEM] {} joined the chat", user);
                                let mut users = online_users_clone.lock().await;
                                users.insert(user_id.to_string(), user.to_string());
                                drop(users);
                                println!("[USERS] Online users: {:?}", online_users_clone.lock().await);
                                println!("[DEBUG] User ID: {}, Username: {}", user_id, user);
                            }
                            Some("leave") => {
                                let user_id = parsed["user_id"].as_str().unwrap_or("unknown");
                                println!("[SYSTEM] {} left the chat", user_id);
                                let mut users = online_users_clone.lock().await;
                                users.remove(user_id);
                                drop(users);
                                println!("[USERS] Online users: {:?}", online_users_clone.lock().await);
                            }
                            Some("ack") => {
                                println!("[SYSTEM] Message acknowledged");
                            }
                            Some("pong") => {
                                println!("[SYSTEM] Pong received");
                            }
                            _ => {
                                println!("[RAW] {}", text);
                            }
                        }
                    } else {
                        println!("[RAW] {}", text);
                    }
                }
                Ok(Message::Close(_)) => {
                    println!("Connection closed by server");
                    break;
                }
                Err(e) => {
                    eprintln!("Error reading message: {}", e);
                    break;
                }
                _ => {}
            }
        }
    });
    
    // Handle user input
    let stdin = tokio::io::stdin();
    let mut reader = BufReader::new(stdin);
    let mut line = String::new();
    
    println!("Available commands:");
    println!("  /broadcast <message> - Send broadcast message");
    println!("  /direct <user_id> <message> - Send direct message");
    println!("  /users - Show online users");
    println!("  /ping - Send ping");
    println!("  /quit - Exit");
    println!();
    println!("Tip: Use /users to see online users and their IDs");
    println!();
    
    loop {
        line.clear();
        print!("> ");
        std::io::Write::flush(&mut std::io::stdout())?;
        
        match reader.read_line(&mut line).await {
            Ok(0) => break, // EOF
            Ok(_) => {
                let input = line.trim();
                
                if input.is_empty() {
                    continue;
                }
                
                if input == "/quit" {
                    break;
                }
                
                if input == "/ping" {
                    let ping_msg = json!({
                        "type": "ping"
                    });
                    write.send(Message::Text(ping_msg.to_string())).await?;
                    continue;
                }
                
                if input == "/users" {
                    let users = online_users.lock().await;
                    if users.is_empty() {
                        println!("[USERS] No other users online");
                    } else {
                        println!("[USERS] Online users:");
                        for (user_id, username) in users.iter() {
                            if *user_id != user_id_clone.to_string() {
                                println!("  {} - {}", user_id, username);
                            }
                        }
                    }
                    drop(users);
                    continue;
                }
                
                if input.starts_with("/broadcast ") {
                    let content = input.strip_prefix("/broadcast ").unwrap_or("");
                    if !content.is_empty() {
                        let broadcast_msg = json!({
                            "type": "broadcast",
                            "content": content,
                            "sender_id": user_id_clone,
                            "timestamp": current_timestamp()
                        });
                        write.send(Message::Text(broadcast_msg.to_string())).await?;
                    }
                    continue;
                }
                
                if input.starts_with("/direct ") {
                    let parts: Vec<&str> = input.splitn(3, ' ').collect();
                    if parts.len() == 3 {
                        let target_id = parts[1];
                        let content = parts[2];
                        
                        let direct_msg = json!({
                            "type": "direct",
                            "content": content,
                            "target_id": target_id,
                            "sender_id": user_id_clone,
                            "timestamp": current_timestamp()
                        });
                        write.send(Message::Text(direct_msg.to_string())).await?;
                    } else {
                        println!("Usage: /direct <user_id> <message>");
                    }
                    continue;
                }
                
                // Default to broadcast
                let broadcast_msg = json!({
                    "type": "broadcast",
                    "content": input,
                    "sender_id": user_id_clone,
                    "timestamp": current_timestamp()
                });
                write.send(Message::Text(broadcast_msg.to_string())).await?;
            }
            Err(e) => {
                eprintln!("Error reading input: {}", e);
                break;
            }
        }
    }
    
    // Wait for read task to finish
    let _ = read_task.await;
    
    println!("Client disconnected");
    Ok(())
}

fn current_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}